# ThirdParty Import
import taichi as Ti
# Local Import
import Config.GlobalParameters as GP

# Taichi Function
@Ti.func
def Sample(SampleTexture, u, v):
	I = Ti.Vector([int(u), int(v)])
	I = max(0, min(GP.FluidsResolution - 1, I))
	
	return SampleTexture[I]

@Ti.func
def Bilerp_1d(vl, vr, frac):
	# frac: [0.0, 1.0]
	return vl + frac * (vr - vl)

@Ti.func
def Bilerp(vf, p):
	u, v = p
	s, t = u - 0.5, v - 0.5
	# floor
	iu, iv = Ti.floor(s), Ti.floor(t)
	# fract
	fu, fv = s - iu, t - iv
	a = Sample(vf, iu, iv)
	b = Sample(vf, iu + 1, iv)
	c = Sample(vf, iu, iv + 1)
	d = Sample(vf, iu + 1, iv + 1)
	return Bilerp_1d(Bilerp_1d(a, b, fu), Bilerp_1d(c, d, fu), fv)

@Ti.func
def MonoCubicCalcBasic(s1, s2, s3, s4, f):

	dk1 = (s3 - s1) / 2
	dk2 = (s4 - s2) / 2
	deltay = s3 - s2

	dk1 = max(0, dk1)
	dk2 = max(0, dk2)
	dkn = dk1 * dk2
	dkb = dkn > 0
	dk1 = dk1 * dkb
	dk2 = dk2 * dkb

	a0 = s2
	a1 = dk1
	a2 = 3 * deltay - 2*dk1 - dk2
	a3 = dk1 + dk2 - 2 * deltay

	result = a3 * (f ** 3) + a2 * (f ** 2) + a1 * f + a0
	return result


@Ti.func
def MonoCubicCalc(vf, iu, iv, du, dv, f):

	s1 = Sample(vf, iu - du, iv - dv)
	s2 = Sample(vf, iu, iv)
	s3 = Sample(vf, iu + du, iv + dv)
	s4 = Sample(vf, iu + 2 * du, iv + 2 * dv)

	return MonoCubicCalcBasic(s1, s2, s3, s4, f)

@Ti.func
def Cubiclerp(vf, p):
	u, v = p
	s, t = u - 0.5, v - 0.5
	iu, iv = Ti.floor(s), Ti.floor(t)
	fu, fv = s - iu, t - iv
	
	f1 = MonoCubicCalc(vf, iu-1, iv, 0, 1, fv)
	f2 = MonoCubicCalc(vf, iu, iv, 0, 1, fv)
	f3 = MonoCubicCalc(vf, iu+1, iv, 0, 1, fv)
	f4 = MonoCubicCalc(vf, iu+2, iv, 0, 1, fv)
	return MonoCubicCalcBasic(f1, f2, f3, f4, fu)

@Ti.kernel
def Divergence(Buffer: Ti.template(), DivBuffer: Ti.template()):
	for i, j in Buffer:
		VelocityL = Sample(Buffer, i - 1, j)
		VelocityR = Sample(Buffer, i + 1, j)
		VelocityB = Sample(Buffer, i, j - 1)
		VelocityT = Sample(Buffer, i, j + 1)
		VelocityC = Sample(Buffer, i, j    )
		if i == 0:
			VelocityL.x = -VelocityC.x
		if i == GP.FluidsResolution - 1:
			VelocityR.x = -VelocityC.x
		if j == 0:
			VelocityB.y = -VelocityC.y
		if j == GP.FluidsResolution - 1:
			VelocityT.y = -VelocityC.y
		DivBuffer[i, j] = (VelocityR.x - VelocityL.x + VelocityT.y - VelocityB.y) * 0.5 * GP.InvDeltaX

@Ti.kernel
def UpSample(InputBuffer: Ti.template(), OutputBuffer: Ti.template(), Scale: Ti.template()):
	for i, j in OutputBuffer:
		Position = (Ti.Vector([i, j]) + 0.5) / Scale
		OutputBuffer[i, j] = Bilerp(InputBuffer, Position)
		# OutputBuffer[i, j] = Sample(InputBuffer, Ti.floor(Position[0]), Ti.floor(Position[1]))

@Ti.kernel
def MulAdd(InputBuffer: Ti.template(), OutputBuffer: Ti.template(), Scale: float, Offset: float):
	for i, j in InputBuffer:
		OutputBuffer[i, j] = InputBuffer[i, j] * Scale + Offset

@Ti.kernel
def FillBuffer(Buffer: Ti.template(), Resolution: float):
	for i, j in Buffer:
		Buffer[i, j] = Ti.Vector([i, j, 0]) / Resolution